from keras.preprocessing.image import ImageDataGenerator
from keras.optimizers import Adam
from keras.datasets import mnist
from keras.utils import np_utils
import kerasModels

BatchSize = 50
keep_prob = 0.5
save_path = 'model/keras_model/cnn_keras.h5'
weight_path = 'model/keras_model/cnn_weights.h5'

#preprocess the train data
datagen = ImageDataGenerator(
    featurewise_center=True,
    featurewise_std_normalization=True,
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    #horizontal_flip=True
    )

# input mnist data
(X_train, Y_train), (X_test, Y_test) = mnist.load_data()

# data pre-processing
X_train = X_train.reshape(-1, 1, 28, 28)/255
X_test = X_test.reshape(-1, 1, 28, 28)/255
Y_train = np_utils.to_categorical(Y_train, num_classes=10)
Y_test = np_utils.to_categorical(Y_test, num_classes=10)

#data pre-processing: normalize
datagen.fit(X_train)

# build the cnn model
model = kerasModels.lenet_5(0.5)
model.summary()

#define optimizer
adam = Adam(lr=3e-4)
model.compile(optimizer=adam, loss='categorical_crossentropy', metrics=['accuracy'])

print('training')
# model.fit(X_train, Y_train, epochs=1, batch_size=BatchSize)
# fits the model on batches with real-time data augmentation:
model.fit_generator(datagen.flow(X_train, Y_train, batch_size=16),
                    steps_per_epoch=len(X_train), epochs=1)

print('testing')
loss, accuracy = model.evaluate(X_test, Y_test)
print('loss: %f, accuracy: %f'%(loss, accuracy))

model.save(save_path)
print('saved model into %s'%save_path)

model.save_weights(weight_path)
print('saved weights into %s'%save_path)


